{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/roy/anaconda3/envs/tf/lib/python3.8/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "2023-07-19 16:13:48.605804: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2023-07-19 16:13:49.294783: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_dataset, load_from_disk\n",
    "from transformers import OpenAIGPTTokenizer\n",
    "import tensorflow as tf\n",
    "from tqdm import tqdm\n",
    "#dataset = load_dataset(\"bookcorpusopen\", split=\"train\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = load_from_disk(\"bookcorpusopen/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "block_size=512"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-07-19 16:13:57.535071: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:995] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2023-07-19 16:13:57.554855: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:995] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2023-07-19 16:13:57.555084: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:995] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n"
     ]
    }
   ],
   "source": [
    "tokenizer = OpenAIGPTTokenizer.from_pretrained('openai-gpt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Map (num_proc=8):   0%|          | 0/10000 [00:00<?, ? examples/s]Token indices sequence length is longer than the specified maximum sequence length for this model (26123 > 512). Running this sequence through the model will result in indexing errors\n",
      "Token indices sequence length is longer than the specified maximum sequence length for this model (41350 > 512). Running this sequence through the model will result in indexing errors\n",
      "Token indices sequence length is longer than the specified maximum sequence length for this model (32885 > 512). Running this sequence through the model will result in indexing errors\n",
      "Token indices sequence length is longer than the specified maximum sequence length for this model (50305 > 512). Running this sequence through the model will result in indexing errors\n",
      "Token indices sequence length is longer than the specified maximum sequence length for this model (108814 > 512). Running this sequence through the model will result in indexing errors\n",
      "Token indices sequence length is longer than the specified maximum sequence length for this model (98069 > 512). Running this sequence through the model will result in indexing errors\n",
      "Token indices sequence length is longer than the specified maximum sequence length for this model (106758 > 512). Running this sequence through the model will result in indexing errors\n",
      "Token indices sequence length is longer than the specified maximum sequence length for this model (154525 > 512). Running this sequence through the model will result in indexing errors\n",
      "                                                                              \r"
     ]
    }
   ],
   "source": [
    "def tokenize_function(examples):\n",
    "    token_ids = [tokenizer(text) for text in examples[\"text\"]]\n",
    "    total_length = [len(t[\"input_ids\"]) for t in token_ids]\n",
    "    total_length = [(l//(block_size+1))*(block_size+1) for l in total_length]\n",
    "    result = []\n",
    " \n",
    "    for i in range(len(total_length)):\n",
    "        result.extend([token_ids[i][\"input_ids\"][j:j+block_size+1] for j in range(0, total_length[i], block_size+1)])\n",
    "    return {\"token_ids\": result}\n",
    " \n",
    "ds_test = ds['train'].select(range(10000))\n",
    " \n",
    "tokenized_datasets = ds_test.map(\n",
    "    tokenize_function, batched=True, num_proc=8, remove_columns=[\"title\", \"text\"], batch_size=100\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                                       \r"
     ]
    }
   ],
   "source": [
    "tokenized_datasets.save_to_disk(\"data/boocorpusopen_10000_513tokens\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1694214/1694214 [08:45<00:00, 3221.76it/s]\n"
     ]
    }
   ],
   "source": [
    "def _int64_feature(value):\n",
    "    \"\"\"Returns an int64_list from a bool / enum / int / uint.\"\"\"\n",
    "    return tf.train.Feature(int64_list=tf.train.Int64List(value=value))\n",
    " \n",
    "def serialize_example(token_ids):\n",
    "    feature = {\n",
    "        'token_ids': _int64_feature(token_ids)\n",
    "    }\n",
    " \n",
    "    example_proto = tf.train.Example(features=tf.train.Features(feature=feature))\n",
    "    return example_proto.SerializeToString()\n",
    " \n",
    "records_num = 100000\n",
    "count = 0\n",
    "for record in tqdm(tokenized_datasets):\n",
    "    if count%records_num == 0:\n",
    "        writer = tf.io.TFRecordWriter(\"dataset/bookcorpus_\"+str(count//records_num)+\".tfrecords\")\n",
    "    writer.write(serialize_example(record['token_ids']))\n",
    "    count += 1\n",
    "    if count%records_num == 0:\n",
    "        writer.close()\n",
    "if writer:\n",
    "    writer.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tf",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.17"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
